/*
 * Copyright 2010-2012 Susanta Tewari. <freecode4susant@users.sourceforge.net>
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */

package coalescent.br;

import coalescent.br.BRSPopGen.Tracker;
import coalescent.data.PopGenData;
import coalescent.model.PopGenModel;

import com.google.common.collect.ImmutableSet;

import commons.core.Statistic;

import java.math.BigDecimal;

/**
 * Provides a skeleton implementation of backward recursion for a statistic defined over population
 * genetic data.
 *
 * @param <T> tracker type
 * @param <A> allele type
 * @param <E> event type
 * @param <M> model type
 * @param <D> data type
 * @param <S> statistic type
 * @param <R> runtime type of this class
 * @author Susanta Tewari
 */
public abstract class BRSPopGen<T extends Tracker<A, E, M, D, S, R>, A, E extends Enum<E>,
                                M extends PopGenModel<E>, D extends PopGenData<A, E, M, D>,
                                S extends Statistic<D, S>,
                                R extends BRSPopGen<T, A, E, M, D, S, R>>
        implements AMCElement<E, R> {

    /**
     * the statistic defines the sample configuration of the recursion
     */
    private S statistic;

    /**
     * Tracker object containing information about the origin of the ancestral configurations
     */
    private T t;

    /**
     * Creates an instance from a given statistic.
     *
     * @param statistic the underlying statistic of this recursion
     */
    protected BRSPopGen(S statistic) {
        this.statistic = statistic;
    }

    /**
     * Returns the underlying statistic of this recursion.
     *
     * @return the underlying statistic
     */
    protected final S getStatistic() {
        return statistic;
    }

    /**
     * Creates a new runtime instance of this recursion from a given statistic. Typically, the
     * implementation would be delegated to the constructor or a factory method of a concrete
     * class. <p>
     * Note that this method is complemented by the method {@code getStatistic()} which
     * returns the underlying statistic. The following steps in order are performed during
     * recursion:
     * <ol>
     * <li>get the statistic of this recursion using {@code getStatistic()}</li>
     * <li>get the mutable version of the statistic using {@code getMutable()}</li>
     * <li>edit the mutable version</li>
     * <li>create an immutable version using {@code getImmutable()}</li>
     * <li>create a recursion using method {@code factory()}</li>
     * </ol>
     *
     * @param statistic statistic
     * @return a new instance of backward recursion based statistic
     */
    protected abstract R factory(S statistic);

    /**
     * Returns the runtime instance of this recursion. Typically, this is implemented in the
     * concrete class by returning {@code this}.
     *
     * @return this instance
     */
    @SuppressWarnings("unchecked")
    protected final R reflect() {
        return (R) this;
    }

    /**
     * Gets the tracker object associated with this ancestral recursion. Trackers are useful during
     * the calculation of transition probability for validation of the ancestral configuration and
     * the calculation itself.
     *
     * @return tracker object associated with this ancestral recursion or {@code null} if this is
     *         the starting object of the recursion
     */
    protected final T getTracker() {
        return t;
    }

    /**
     * A new tracker object is associated with a new ancestral configuration created in
     * {@code applyEvent()}. Setter is used to avoid protected field.
     *
     * @param t a newly create tracker object
     */
    protected final void setTracker(T t) {
        this.t = t;
    }

    /**
     * The probability of observing next an event of the specified type, when {@code n}
     * ancestors are present.
     *
     * @param type a supported event type
     * @param ancestralConfig
     * @return the event probability
     */
    protected abstract BigDecimal eventProb(E type, R ancestralConfig);

    /**
     * Conditional probability that this recursion configuration generated the given ancestral
     * configuration for the specified event type.
     *
     * @param type event type of MC step
     * @param ancestralConfig recursion ancestral configuration
     * @return the conditional probability
     */
    protected abstract BigDecimal cndAncestralProb(E type, R ancestralConfig);

    /**
     * Run checks to ensure that the specified configuration was generated by this recursion.
     * Subclasses can add additional checks.
     *
     * @param type event type that created the ancestral configuration
     * @param ancestralConfig an ancestral configuration supposed to be created by this recursion
     */
    protected void checkForCndAncestralProb(E type, R ancestralConfig) {

        if (type == null) {
            throw new NullPointerException("type==null");
        }

        if (ancestralConfig == null) {
            throw new NullPointerException("ancestralConfig==null");
        }


        // check producer: this recursion should be the ancestor`s producer
        if (this != getTracker().getProducer()) {
            throw new IllegalArgumentException("invalid ancestral config: producer check failed");
        }


        // check event type: argument event type must match with that of the ancestor`s
        if (type != ancestralConfig.getTracker().getEventType()) {
            throw new IllegalArgumentException("invalid ancestral config: event type check failed");
        }


        // check allele: the allele in ancestor must be present in this recursion
        if (!getStatistic().getData().getAlleles().contains(
                ancestralConfig.getTracker().getAllele())) {
            throw new IllegalArgumentException("invalid ancestral config: allele check failed");
        }
    }

    /**
     * Returns the event types of the underlying population genetic model.
     *
     * @return model event types
     */
    @Override
    public final ImmutableSet<E> transitionTypes() {
        return getStatistic().getData().getModel().getEventTypes();
    }

    /**
     * Method description
     *
     *
     * @param type
     * @param ancestralConfig
     *
     * @return
     */
    @Override
    public final BigDecimal transitionProb(E type, R ancestralConfig) {

        checkForCndAncestralProb(type, ancestralConfig);

        BigDecimal eventProb        = eventProb(type, ancestralConfig);
        BigDecimal cndAncestralProb = cndAncestralProb(type, ancestralConfig);

        return eventProb.multiply(cndAncestralProb);
    }

    /**
     * Default implementation: {@code alleleCount + sampleSize - 1}
     *
     * @return cut-set index
     */
    @Override
    public int transitionIndex() {

        D data          = getStatistic().getData();
        int sampleSize  = data.getSampleSize();
        int alleleCount = data.getAlleleCount();

        return alleleCount + sampleSize - 1;
    }

    /**
     * Recursion reaches boundary when the underlying configuration is MRCA.
     *
     * @return {@code true} if the underlying configuration is MRCA; {@code false} otherwise
     */
    @Override
    public boolean isSink() {
        return getStatistic().getData().isMRCA();
    }

    /**
     * Only one boundary is possible; when the underlying configuration reaches MRCA. Its
     * probability is 1.
     *
     * @return
     */
    @Override
    public BigDecimal sinkProbability() {

        if (!isSink()) {
            throw new IllegalArgumentException("recursion is not a boundary");
        }

        return BigDecimal.valueOf(1.0);
    }

    /**
     * Tracks information on ancestral configurations created by this recursion. Such information is
     * useful in computing transition probabilities. Tracker objects are created to associate with
     * the new ancestral configurations created in {@code applyEvent()}.
     *
     * @param <A> allele type
     * @param <E> event type
     * @param <M> model type
     * @param <D> data type
     * @param <S> statistic type
     * @param <R> runtime type of this class
     */
    @SuppressWarnings("ProtectedInnerClass")
    protected static class Tracker<A, E extends Enum<E>, M extends PopGenModel<E>,
                                   D extends PopGenData<A, E, M, D>, S extends Statistic<D, S>,
                                   R extends BRSPopGen<?, A, E, M, D, S, R>> {

        /** Field description */
        private E eventType;

        /** Field description */
        private R producer;

        /** Field description */
        private A allele;

        /**
         * Constructs ...
         *
         *
         * @param allele
         * @param eventType
         * @param producer
         */
        protected Tracker(A allele, E eventType, R producer) {

            this.eventType = eventType;
            this.producer  = producer;
            this.allele    = allele;
        }

        /**
         * Method description
         *
         *
         * @return
         */
        protected final E getEventType() {
            return eventType;
        }

        /**
         * Method description
         *
         *
         * @return
         */
        protected final R getProducer() {
            return producer;
        }

        /**
         * Method description
         *
         *
         * @return
         */
        protected final A getAllele() {
            return allele;
        }
    }
}
